# encoding=utf-8
# 使用的是python2(inception-v4)环境
import torch
import imp
import torch.onnx
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def test():
    MainModel = imp.load_source('MainModel', "inception_v4.py")
    model = torch.load("inception-v4.pth")
    model.eval()
    pthfile = 'inception-v4.pth'
    loaded_model = torch.load(pthfile, map_location=device)
    model.load_state_dict(loaded_model.state_dict())
    model = model.to(device)
    model.eval()

    # data type nchw
    dummy_input1 = torch.randn(1, 3, 299, 299).to(device)
    input_names = ["input"]
    output_names = ["output"]
    torch.onnx.export(model,
                      (dummy_input1,),
                      "inception-v4.onnx",
                      verbose=True,
                      input_names=input_names,
                      output_names=output_names)

# AttributeError: 'long' object has no attribute 'item' 报错处理, 换位python2解释器


if __name__ == "__main__":
    test()